import random

import numpy as np

import mnist_loader


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


def sigmoid_prime(z):
    return sigmoid(z) * (1 - sigmoid(z))


class Network:
    def __init__(self, sizes):
        # size is the nodes number im every layer, like [2, 3, 1]
        self.num_layers = len(sizes)
        self.sizes = sizes
        self.bias = [np.random.randn(y, 1) for y in sizes[1:]]
        self.weights = [np.random.randn(y, x) for x, y in zip(sizes[:-1], sizes[1:])]

    def feedforward(self, a):
        for b, w in zip(self.bias, self.weights):
            a = sigmoid(np.dot(w, a) + b)
        return a

    def SGD(self, tr, epochs, batch_size, learn_rate, te):
        if te:
            n_test = len(te[0])
        train_data_len = len(tr[0])
        for j in range(epochs):
            # random.shuffle(tr)
            batchs = [[(tr[0][i], tr[1][i]) for i in range(k, k + batch_size)] for k in range(0, train_data_len, batch_size)]
            for one_batch in batchs:
                self.update_batch(one_batch, learn_rate)
            if te:
                print('epoch {0}:{1} {2}'.format(j, self.evaluate(te), n_test))
            else:
                print('epoch {0} completed'.format(j))

    def evaluate(self, test_data):
        te_zip = [(test_data[0][i], test_data[1][i]) for i in range(len(test_data[0]))]
        test_results = [(np.argmax(self.feedforward(x)), y) for x, y in te_zip]
        return sum(int(x == y) for x, y in test_results)

    def update_batch(self, one_batch, learn_rate):
        delta_b = [np.zeros(b.shape) for b in self.bias]
        delta_w = [np.zeros(w.shape) for w in self.weights]
        for x, y in one_batch:
            # method backprop is to calculate derivative of loss for weight and bias
            delta_b_sample, delta_w_sample = self.backprop(x, y)
            delta_b = [db + dnb for db, dnb in zip(delta_b, delta_b_sample)]
            delta_w = [dw + dnw for dw, dnw in zip(delta_w, delta_w_sample)]
        self.weights = [w - (learn_rate/len(one_batch)) * nw for w, nw in zip(self.weights, delta_w)]
        self.bias = [b - (learn_rate/len(one_batch)) * nb for b, nb in zip(self.bias, delta_b)]

    def backprop(self, x, y):
        nabla_b = [np.zeros(b.shape) for b in self.bias]
        nabla_w = [np.zeros(w.shape) for w in self.weights]

        # feedforward operation
        activation = x
        activations = [x]
        zs = []
        for b, w in zip(self.bias, self.weights):
            z = np.dot(w, activation) + b
            zs.append(z)
            activation = sigmoid(z)
            activations.append(activation)

        # back forward operation
        delta = self.cost_derivative(activations[-1], y) * sigmoid_prime(zs[-1])
        nabla_b[-1] = delta
        nabla_w[-1] = np.dot(delta, activations[-2].transpose())
        for layer_number in range(2, self.num_layers):
            z = zs[-layer_number]
            sp = sigmoid_prime(z)
            delta = np.dot(self.weights[-layer_number + 1].transpose(), delta) * sp
            nabla_b[-layer_number] = delta
            nabla_w[-layer_number] = np.dot(delta, activations[-layer_number - 1].transpose())
        return nabla_b, nabla_w

    def cost_derivative(self, output_activations, y):
        return output_activations - y


def main():
    tr, va, te = mnist_loader.load_data_wrapper()
    n = Network([784, 30, 10])
    n.SGD(tr, 30, 10, 3, te=te)


if __name__ == '__main__':
    main()
